package com.xw.csrf;

import cn.hutool.core.util.RandomUtil;
import cn.hutool.crypto.digest.DigestUtil;
import org.springframework.stereotype.Component;

import javax.servlet.http.Cookie;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.HashMap;
import java.util.Map;
/**
 * @Description: CsrfTokenManager
 * @Author: ASUS
 * @Date: 2024/12/19 09:15
 * @Version: 1.0
 */
@Component
public class CsrfTokenManager {

    private static final String CSRF_TOKEN_KEY = "_csrf_token";
    private static final String CSRF_TOKEN_HEADER = "X-CSRF-Token";
    private Map<String, String> tokenStore = new HashMap<>();

    // 生成CSRF Token
    public String generateToken() {
        String token = RandomUtil.randomString(32);
        return DigestUtil.md5Hex(token);
    }

    // 将Token存储到内存（简单示例，实际可考虑更可靠的存储如缓存等）
    public void storeToken(String tokenId, String token) {
        tokenStore.put(tokenId, token);
    }

    // 从内存中获取Token
    public String getToken(String tokenId) {
        return tokenStore.get(tokenId);
    }

    // 从请求中获取Token（先从请求头获取，没有则尝试从Cookie获取）
    public String getTokenFromRequest(HttpServletRequest request) {
        String token = request.getHeader(CSRF_TOKEN_HEADER);
        if (token == null) {
            Cookie[] cookies = request.getCookies();
            if (cookies!= null) {
                for (Cookie cookie : cookies) {
                    if (CSRF_TOKEN_KEY.equals(cookie.getName())) {
                        token = cookie.getValue();
                        break;
                    }
                }
            }
        }
        return token;
    }

    // 在响应中设置Token到Cookie（方便前端后续获取）
    public void setTokenInCookie(HttpServletResponse response, String token) {
        Cookie cookie = new Cookie(CSRF_TOKEN_KEY, token);
        cookie.setHttpOnly(true);
        cookie.setPath("/");
        response.addCookie(cookie);
    }

    // 验证Token是否有效
    public boolean validateToken(String tokenId, String token) {
        String storedToken = getToken(tokenId);
        return storedToken!= null && storedToken.equals(token);
    }
}